#include "number_classifier.h"
#include <iostream>

NumberClassifier::NumberClassifier() : model_loaded_(false) {
    svm_model_ = cv::ml::SVM::create();
}

bool NumberClassifier::loadModel(const std::string& model_path) {
    try {
        cv::Ptr<cv::ml::SVM> loaded = cv::ml::SVM::load(model_path);
        if (loaded.empty()) {
            std::cout << "Failed to load SVM model: Model is empty or not found: " << model_path << std::endl;
            model_loaded_ = false;
            return false;
        }

        svm_model_ = loaded;
        model_loaded_ = true;
        std::cout << "SVM model loaded successfully from: " << model_path << std::endl;
        return true;
    } catch (const cv::Exception& e) {
        std::cout << "Failed to load SVM model: " << e.what() << std::endl;
        model_loaded_ = false;
        return false;
    }
}

cv::Mat NumberClassifier::preprocessROI(const cv::Mat& roi) {
    cv::Mat processed;
    if (roi.empty()) return cv::Mat();

    // 转换为灰度图（兼容单/三通道输入）
    if (roi.channels() == 3) {
        cv::cvtColor(roi, processed, cv::COLOR_BGR2GRAY);
    } else {
        roi.copyTo(processed);
    }
    if (processed.empty()) return cv::Mat();

    // preserve aspect ratio: pad to square then resize
    int w = processed.cols;
    int h = processed.rows;
    int max_dim = std::max(w, h);
    cv::Mat square = cv::Mat::zeros(max_dim, max_dim, processed.type());
    int x = (max_dim - w) / 2;
    int y = (max_dim - h) / 2;
    processed.copyTo(square(cv::Rect(x, y, w, h)));

    cv::resize(square, processed, cv::Size(img_size_, img_size_));

    // 二值化
    cv::threshold(processed, processed, 0, 255, cv::THRESH_BINARY | cv::THRESH_OTSU);
    
    // 转换为浮点数并归一化
    processed.convertTo(processed, CV_32F, 1.0 / 255.0);
    
    return processed.reshape(1, 1); // 展平为一行
}

int NumberClassifier::predict(const cv::Mat& roi) {
    if (!model_loaded_) {
        std::cout << "Model not loaded!" << std::endl;
        return -1;
    }
    
    try {
        cv::Mat features = preprocessROI(roi);
        float result = svm_model_->predict(features);
        return static_cast<int>(result);
    } catch (const cv::Exception& e) {
        std::cout << "Prediction error: " << e.what() << std::endl;
        return -1;
    }
}